

<!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  
  <title>mindspore.nn.layer.embedding &mdash; MindSpore master documentation</title>
  

  
  <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />

  
  

  
  

  

  
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
    
      <script type="text/javascript" id="documentation_options" data-url_root="../../../../" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
        <script src="../../../../_static/language_data.js"></script>
        <script async="async" src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/latest.js?config=TeX-AMS-MML_HTMLorMML"></script>
    
    <script type="text/javascript" src="../../../../_static/js/theme.js"></script>

    
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav">

   
  <div class="wy-grid-for-nav">
    
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
          

          
            <a href="../../../../index.html" class="icon icon-home"> MindSpore
          

          
          </a>

          
            
            
          

          
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

          
        </div>

        
        <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          
            
            
              
            
            
              <p class="caption"><span class="caption-text">MindSpore Python API</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.html">mindspore</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.common.initializer.html">mindspore.common.initializer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.communication.html">mindspore.communication</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.compression.html">mindspore.compression</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.context.html">mindspore.context</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.html">mindspore.dataset</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.audio.html">mindspore.dataset.audio</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.config.html">mindspore.dataset.config</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.text.html">mindspore.dataset.text</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.transforms.html">mindspore.dataset.transforms</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.vision.html">mindspore.dataset.vision</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.mindrecord.html">mindspore.mindrecord</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.html">mindspore.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.probability.html">mindspore.nn.probability</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.transformer.html">mindspore.nn.transformer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.numpy.html">mindspore.numpy</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.ops.html">mindspore.ops</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.parallel.html">mindspore.parallel</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.parallel.nn.html">mindspore.parallel.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.profiler.html">mindspore.profiler</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.scipy.html">mindspore.scipy</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.train.html">mindspore.train</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.boost.html">mindspore.boost</a></li>
</ul>
<p class="caption"><span class="caption-text">MindSpore C++ API</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://www.mindspore.cn/lite/api/zh-CN/master/api_cpp/mindspore.html">MindSpore Lite↗</a></li>
</ul>

            
          
        </div>
        
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">

      
      <nav class="wy-nav-top" aria-label="top navigation">
        
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">MindSpore</a>
        
      </nav>


      <div class="wy-nav-content">
        
        <div class="rst-content">
        
          

















<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="wy-breadcrumbs">
    
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
        
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
        
      <li>mindspore.nn.layer.embedding</li>
    
    
      <li class="wy-breadcrumbs-aside">
        
      </li>
    
  </ul>

  
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
            
  <h1>Source code for mindspore.nn.layer.embedding</h1><div class="highlight"><pre>
<span></span><span class="c1"># Copyright 2020-2021 Huawei Technologies Co., Ltd</span>
<span class="c1">#</span>
<span class="c1"># Licensed under the Apache License, Version 2.0 (the &quot;License&quot;);</span>
<span class="c1"># you may not use this file except in compliance with the License.</span>
<span class="c1"># You may obtain a copy of the License at</span>
<span class="c1">#</span>
<span class="c1"># http://www.apache.org/licenses/LICENSE-2.0</span>
<span class="c1">#</span>
<span class="c1"># Unless required by applicable law or agreed to in writing, software</span>
<span class="c1"># distributed under the License is distributed on an &quot;AS IS&quot; BASIS,</span>
<span class="c1"># WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.</span>
<span class="c1"># See the License for the specific language governing permissions and</span>
<span class="c1"># limitations under the License.</span>
<span class="c1"># ============================================================================</span>
<span class="sd">&quot;&quot;&quot;embedding&quot;&quot;&quot;</span>
<span class="kn">import</span> <span class="nn">mindspore.common.dtype</span> <span class="k">as</span> <span class="nn">mstype</span>
<span class="kn">from</span> <span class="nn">mindspore</span> <span class="kn">import</span> <span class="n">log</span> <span class="k">as</span> <span class="n">logger</span>
<span class="kn">from</span> <span class="nn">mindspore.common.tensor</span> <span class="kn">import</span> <span class="n">Tensor</span>
<span class="kn">from</span> <span class="nn">mindspore.ops</span> <span class="kn">import</span> <span class="n">operations</span> <span class="k">as</span> <span class="n">P</span>
<span class="kn">from</span> <span class="nn">mindspore.ops</span> <span class="kn">import</span> <span class="n">functional</span> <span class="k">as</span> <span class="n">F</span>
<span class="kn">from</span> <span class="nn">mindspore.common.parameter</span> <span class="kn">import</span> <span class="n">Parameter</span>
<span class="kn">from</span> <span class="nn">mindspore.common.initializer</span> <span class="kn">import</span> <span class="n">initializer</span>
<span class="kn">from</span> <span class="nn">mindspore.communication.management</span> <span class="kn">import</span> <span class="n">get_group_size</span><span class="p">,</span> <span class="n">get_rank</span>
<span class="kn">from</span> <span class="nn">mindspore.context</span> <span class="kn">import</span> <span class="n">ParallelMode</span>
<span class="kn">from</span> <span class="nn">mindspore.parallel._utils</span> <span class="kn">import</span> <span class="n">_get_parallel_mode</span><span class="p">,</span> <span class="n">_get_full_batch</span>
<span class="kn">from</span> <span class="nn">mindspore.parallel._ps_context</span> <span class="kn">import</span> <span class="n">_is_role_worker</span><span class="p">,</span> <span class="n">_get_ps_context</span>
<span class="kn">from</span> <span class="nn">mindspore.parallel._ps_context</span> <span class="kn">import</span> <span class="n">_insert_hash_table_size</span><span class="p">,</span> <span class="n">_set_cache_enable</span><span class="p">,</span> <span class="n">_set_rank_id</span>
<span class="kn">from</span> <span class="nn">mindspore</span> <span class="kn">import</span> <span class="n">context</span>
<span class="kn">from</span> <span class="nn">mindspore._checkparam</span> <span class="kn">import</span> <span class="n">Rel</span>
<span class="kn">from</span> <span class="nn">mindspore._checkparam</span> <span class="kn">import</span> <span class="n">Validator</span> <span class="k">as</span> <span class="n">validator</span>
<span class="kn">from</span> <span class="nn">mindspore.ops.primitive</span> <span class="kn">import</span> <span class="n">constexpr</span>
<span class="kn">from</span> <span class="nn">.basic</span> <span class="kn">import</span> <span class="n">ClipByNorm</span>
<span class="kn">from</span> <span class="nn">.math</span> <span class="kn">import</span> <span class="n">Range</span>
<span class="kn">from</span> <span class="nn">..cell</span> <span class="kn">import</span> <span class="n">Cell</span>

<span class="n">__all__</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;Embedding&#39;</span><span class="p">,</span> <span class="s1">&#39;EmbeddingLookup&#39;</span><span class="p">,</span> <span class="s1">&#39;MultiFieldEmbeddingLookup&#39;</span><span class="p">]</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_check_input_2d</span><span class="p">(</span><span class="n">input_shape</span><span class="p">,</span> <span class="n">param_name</span><span class="p">,</span> <span class="n">func_name</span><span class="p">):</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">input_shape</span><span class="p">)</span> <span class="o">!=</span> <span class="mi">2</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="n">func_name</span><span class="si">}</span><span class="s2">&#39;, the dimension of &#39;</span><span class="si">{</span><span class="n">param_name</span><span class="si">}</span><span class="s2">&#39; should be 2d, but got </span><span class="si">{</span><span class="nb">len</span><span class="p">(</span><span class="n">input_shape</span><span class="p">)</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
    <span class="k">return</span> <span class="kc">True</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_check_input_dtype</span><span class="p">(</span><span class="n">input_dtype</span><span class="p">,</span> <span class="n">param_name</span><span class="p">,</span> <span class="n">allow_dtypes</span><span class="p">,</span> <span class="n">cls_name</span><span class="p">):</span>
    <span class="n">validator</span><span class="o">.</span><span class="n">check_type_name</span><span class="p">(</span><span class="n">param_name</span><span class="p">,</span> <span class="n">input_dtype</span><span class="p">,</span> <span class="n">allow_dtypes</span><span class="p">,</span> <span class="n">cls_name</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">Embedding</span><span class="p">(</span><span class="n">Cell</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A simple lookup table that stores embeddings of a fixed dictionary and size.</span>

<span class="sd">    This module is often used to store word embeddings and retrieve them using</span>
<span class="sd">    indices. The input to the module is a list of indices, and the output is</span>
<span class="sd">    the corresponding word embeddings.</span>

<span class="sd">    Note:</span>
<span class="sd">        When &#39;use_one_hot&#39; is set to True, the type of the `x` must be mindspore.int32.</span>

<span class="sd">    Args:</span>
<span class="sd">        vocab_size (int): Size of the dictionary of embeddings.</span>
<span class="sd">        embedding_size (int): The size of each embedding vector.</span>
<span class="sd">        use_one_hot (bool): Specifies whether to apply one_hot encoding form. Default: False.</span>
<span class="sd">        embedding_table (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the embedding_table.</span>
<span class="sd">            Refer to class `initializer` for the values of string when a string</span>
<span class="sd">            is specified. Default: &#39;normal&#39;.</span>
<span class="sd">        dtype (:class:`mindspore.dtype`): Data type of `x`. Default: mindspore.float32.</span>
<span class="sd">        padding_idx (int, None): When the padding_idx encounters index, the output embedding vector of this index</span>
<span class="sd">                                 will be initialized to zero. Default: None. The feature is inactivated.</span>
<span class="sd">    Inputs:</span>
<span class="sd">        - **x** (Tensor) - Tensor of shape :math:`(\text{batch_size}, \text{x_length})`. The elements of</span>
<span class="sd">          the Tensor must be integer and not larger than vocab_size. Otherwise the corresponding embedding vector will</span>
<span class="sd">          be zero. The data type is int32 or int64.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor of shape :math:`(\text{batch_size}, \text{x_length}, \text{embedding_size})`.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If `vocab_size` or `embedding_size` is not an int.</span>
<span class="sd">        TypeError: If `use_one_hot` is not a bool.</span>
<span class="sd">        ValueError: If `padding_idx` is an int which not in range [0, `vocab_size`].</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; net = nn.Embedding(20000, 768,  True)</span>
<span class="sd">        &gt;&gt;&gt; x = Tensor(np.ones([8, 128]), mindspore.int32)</span>
<span class="sd">        &gt;&gt;&gt; # Maps the input word IDs to word embedding.</span>
<span class="sd">        &gt;&gt;&gt; output = net(x)</span>
<span class="sd">        &gt;&gt;&gt; result = output.shape</span>
<span class="sd">        &gt;&gt;&gt; print(result)</span>
<span class="sd">        (8, 128, 768)</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">vocab_size</span><span class="p">,</span> <span class="n">embedding_size</span><span class="p">,</span> <span class="n">use_one_hot</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">embedding_table</span><span class="o">=</span><span class="s1">&#39;normal&#39;</span><span class="p">,</span>
                 <span class="n">dtype</span><span class="o">=</span><span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">,</span> <span class="n">padding_idx</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize Embedding.&quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">Embedding</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">vocab_size</span> <span class="o">=</span> <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;vocab_size&#39;</span><span class="p">,</span> <span class="n">vocab_size</span><span class="p">,</span> <span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">embedding_size</span> <span class="o">=</span> <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;embedding_size&#39;</span><span class="p">,</span> <span class="n">embedding_size</span><span class="p">,</span> <span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;use_one_hot&#39;</span><span class="p">,</span> <span class="n">use_one_hot</span><span class="p">,</span> <span class="p">[</span><span class="nb">bool</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_subclass</span><span class="p">(</span><span class="s2">&quot;dtype&quot;</span><span class="p">,</span> <span class="n">dtype</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">number_type</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">use_one_hot</span> <span class="o">=</span> <span class="n">use_one_hot</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span> <span class="o">=</span> <span class="n">dtype</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">init_tensor</span> <span class="o">=</span> <span class="n">initializer</span><span class="p">(</span><span class="n">embedding_table</span><span class="p">,</span> <span class="p">[</span><span class="n">vocab_size</span><span class="p">,</span> <span class="n">embedding_size</span><span class="p">])</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">padding_idx</span> <span class="o">=</span> <span class="n">padding_idx</span>
        <span class="k">if</span> <span class="n">padding_idx</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">padding_idx</span> <span class="o">=</span> <span class="n">validator</span><span class="o">.</span><span class="n">check_int_range</span><span class="p">(</span><span class="n">padding_idx</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">vocab_size</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">INC_BOTH</span><span class="p">,</span>
                                                         <span class="s2">&quot;padding_idx&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">init_tensor</span><span class="p">,</span> <span class="n">Tensor</span><span class="p">)</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">init_tensor</span><span class="o">.</span><span class="n">init</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">init_tensor</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">init_tensor</span><span class="o">.</span><span class="n">init_data</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">init_tensor</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">init_tensor</span><span class="o">.</span><span class="n">asnumpy</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">init_tensor</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">padding_idx</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">init_tensor</span> <span class="o">=</span> <span class="n">Tensor</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">init_tensor</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">embedding_table</span> <span class="o">=</span> <span class="n">Parameter</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">init_tensor</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="s1">&#39;embedding_table&#39;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">expand</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">ExpandDims</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">reshape_flat</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Reshape</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">shp_flat</span> <span class="o">=</span> <span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">gather</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Gather</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">one_hot</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">OneHot</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">on_value</span> <span class="o">=</span> <span class="n">Tensor</span><span class="p">(</span><span class="mf">1.0</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">off_value</span> <span class="o">=</span> <span class="n">Tensor</span><span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">array_mul</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">MatMul</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">reshape</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Reshape</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">get_shp</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Shape</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">construct</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ids</span><span class="p">):</span>
        <span class="n">extended_ids</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">expand</span><span class="p">(</span><span class="n">ids</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>
        <span class="n">out_shape</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_shp</span><span class="p">(</span><span class="n">ids</span><span class="p">)</span> <span class="o">+</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">embedding_size</span><span class="p">,)</span>
        <span class="n">flat_ids</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshape_flat</span><span class="p">(</span><span class="n">extended_ids</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">shp_flat</span><span class="p">)</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">use_one_hot</span><span class="p">:</span>
            <span class="n">one_hot_ids</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">one_hot</span><span class="p">(</span><span class="n">flat_ids</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">vocab_size</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">on_value</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">off_value</span><span class="p">)</span>
            <span class="n">output_for_reshape</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">array_mul</span><span class="p">(</span><span class="n">one_hot_ids</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">embedding_table</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">output_for_reshape</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">gather</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">embedding_table</span><span class="p">,</span> <span class="n">flat_ids</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>

        <span class="n">output</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">output_for_reshape</span><span class="p">,</span> <span class="n">out_shape</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">output</span>

    <span class="k">def</span> <span class="nf">extend_repr</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="n">s</span> <span class="o">=</span> <span class="s1">&#39;vocab_size=</span><span class="si">{}</span><span class="s1">, embedding_size=</span><span class="si">{}</span><span class="s1">, use_one_hot=</span><span class="si">{}</span><span class="s1">, embedding_table=</span><span class="si">{}</span><span class="s1">, dtype=</span><span class="si">{}</span><span class="s1">, padding_idx=</span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">vocab_size</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">embedding_size</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">use_one_hot</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">embedding_table</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">padding_idx</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">s</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_make_axis_range</span><span class="p">(</span><span class="n">start</span><span class="p">,</span> <span class="n">end</span><span class="p">):</span>
    <span class="n">axis</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="n">start</span><span class="p">,</span> <span class="n">end</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">axis</span>


<span class="k">class</span> <span class="nc">EmbeddingLookup</span><span class="p">(</span><span class="n">Cell</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Returns a slice of the input tensor based on the specified indices.</span>

<span class="sd">    Note:</span>
<span class="sd">        When &#39;target&#39; is set to &#39;CPU&#39;, this module will use</span>
<span class="sd">        P.EmbeddingLookup().add_prim_attr(&#39;primitive_target&#39;, &#39;CPU&#39;) which</span>
<span class="sd">        specified &#39;offset = 0&#39; to lookup table.</span>
<span class="sd">        When &#39;target&#39; is set to &#39;DEVICE&#39;, this module will use P.Gather() which</span>
<span class="sd">        specified &#39;axis = 0&#39; to lookup table.</span>
<span class="sd">        In field slice mode, the manual_shapes must be given. It is a tuple ,where</span>
<span class="sd">        the element is vocab[i], vocab[i] is the row numbers for i-th part.</span>

<span class="sd">    Args:</span>
<span class="sd">        vocab_size (int): Size of the dictionary of embeddings.</span>
<span class="sd">        embedding_size (int): The size of each embedding vector.</span>
<span class="sd">        param_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the embedding_table.</span>
<span class="sd">            Refer to class `initializer` for the values of string when a string</span>
<span class="sd">            is specified. Default: &#39;normal&#39;.</span>
<span class="sd">        target (str): Specifies the target where the op is executed. The value must in</span>
<span class="sd">            [&#39;DEVICE&#39;, &#39;CPU&#39;]. Default: &#39;CPU&#39;.</span>
<span class="sd">        slice_mode (str): The slicing way in semi_auto_parallel/auto_parallel. The value must get through</span>
<span class="sd">            :class:`mindspore.nn.EmbeddingLookup`. Default: :class:`mindspore.nn.EmbeddingLookup.BATCH_SLICE`.</span>
<span class="sd">        manual_shapes (tuple): The accompaniment array in field slice mode.</span>
<span class="sd">        max_norm (Union[float, None]): A maximum clipping value. The data type must be float16, float32</span>
<span class="sd">                                       or None. Default: None</span>
<span class="sd">        sparse (bool): Using sparse mode. When &#39;target&#39; is set to &#39;CPU&#39;, &#39;sparse&#39; has to be true. Default: True.</span>
<span class="sd">        vocab_cache_size (int): Cache size of the dictionary of embeddings. Default: 0. It is valid only in</span>
<span class="sd">            parameter server trainning mode and &#39;DEVICE&#39; target. And the moment parameter of corresponding</span>
<span class="sd">            optimizer will also be set to the cache size. In addition, it should be noted that it will cost the &#39;DEVICE&#39;</span>
<span class="sd">            memory, so suggests setting a reasonable value to avoid insufficient memory.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **input_indices** (Tensor) - The shape of tensor is :math:`(y_1, y_2, ..., y_S)`.</span>
<span class="sd">          Specifies the indices of elements of the original Tensor. Values can be out of range of embedding_table,</span>
<span class="sd">          and the exceeding part will be filled with 0 in the output. Values does not support negative and the result</span>
<span class="sd">          is undefined if values are negative. Input_indices must only be a 2d tensor in</span>
<span class="sd">          this interface when run in semi auto parallel/auto parallel mode.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor, the shape of tensor is :math:`(z_1, z_2, ..., z_N)`.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If `vocab_size` or `embedding_size` or `vocab_cache_size` is not an int.</span>
<span class="sd">        TypeError: If `sparse` is not a bool or `manual_shapes` is not a tuple.</span>
<span class="sd">        ValueError: If `vocab_size` or `embedding_size` is less than 1.</span>
<span class="sd">        ValueError: If `vocab_cache_size` is less than 0.</span>
<span class="sd">        ValueError: If `target` is neither &#39;CPU&#39; nor &#39;DEVICE&#39;.</span>
<span class="sd">        ValueError: If `slice_mode` is not one of &#39;batch_slice&#39; or &#39;field_slice&#39; or</span>
<span class="sd">                    &#39;table_row_slice&#39; or &#39;table_column_slice&#39;.</span>
<span class="sd">        ValueError: If `sparse` is False and `target` is &#39;CPU&#39;.</span>
<span class="sd">        ValueError: If `slice_mode` is &#39;field_slice&#39; and `manual_shapes` is None.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; input_indices = Tensor(np.array([[1, 0], [3, 2]]), mindspore.int32)</span>
<span class="sd">        &gt;&gt;&gt; result = nn.EmbeddingLookup(4,2)(input_indices)</span>
<span class="sd">        &gt;&gt;&gt; print(result.shape)</span>
<span class="sd">        (2, 2, 2)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">BATCH_SLICE</span> <span class="o">=</span> <span class="s2">&quot;batch_slice&quot;</span>
    <span class="n">FIELD_SLICE</span> <span class="o">=</span> <span class="s2">&quot;field_slice&quot;</span>
    <span class="n">TABLE_ROW_SLICE</span> <span class="o">=</span> <span class="s2">&quot;table_row_slice&quot;</span>
    <span class="n">TABLE_COLUMN_SLICE</span> <span class="o">=</span> <span class="s2">&quot;table_column_slice&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">vocab_size</span><span class="p">,</span> <span class="n">embedding_size</span><span class="p">,</span> <span class="n">param_init</span><span class="o">=</span><span class="s1">&#39;normal&#39;</span><span class="p">,</span>
                 <span class="n">target</span><span class="o">=</span><span class="s1">&#39;CPU&#39;</span><span class="p">,</span> <span class="n">slice_mode</span><span class="o">=</span><span class="s1">&#39;batch_slice&#39;</span><span class="p">,</span> <span class="n">manual_shapes</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                 <span class="n">max_norm</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">sparse</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">vocab_cache_size</span><span class="o">=</span><span class="mi">0</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize EmbeddingLookup.&quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">EmbeddingLookup</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;sparse&#39;</span><span class="p">,</span> <span class="n">sparse</span><span class="p">,</span> <span class="p">[</span><span class="nb">bool</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">vocab_size</span> <span class="o">=</span> <span class="n">validator</span><span class="o">.</span><span class="n">check_positive_int</span><span class="p">(</span><span class="n">vocab_size</span><span class="p">,</span> <span class="s1">&#39;vocab_size&#39;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">vocab_cache_size</span> <span class="o">=</span> <span class="n">validator</span><span class="o">.</span><span class="n">check_non_negative_int</span><span class="p">(</span><span class="n">vocab_cache_size</span><span class="p">,</span> <span class="s1">&#39;vocab_cache_size&#39;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">target</span> <span class="o">=</span> <span class="n">target</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">sparse</span> <span class="o">=</span> <span class="n">sparse</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cache_enable</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">vocab_cache_size</span> <span class="o">&gt;</span> <span class="mi">0</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">forward_unique</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_string</span><span class="p">(</span><span class="n">target</span><span class="p">,</span> <span class="p">[</span><span class="s1">&#39;CPU&#39;</span><span class="p">,</span> <span class="s1">&#39;DEVICE&#39;</span><span class="p">],</span> <span class="s1">&#39;target&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">sparse</span> <span class="ow">and</span> <span class="n">target</span> <span class="o">==</span> <span class="s1">&#39;CPU&#39;</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="si">}</span><span class="s2">&#39;, &#39;sparse&#39; must be True when &#39;target&#39; is </span><span class="se">\&quot;</span><span class="s2">CPU</span><span class="se">\&quot;</span><span class="s2">, &quot;</span>
                             <span class="sa">f</span><span class="s2">&quot;but got &#39;sparse&#39;: </span><span class="si">{</span><span class="n">sparse</span><span class="si">}</span><span class="s2"> and &#39;target&#39;: </span><span class="si">{</span><span class="n">target</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">sparse</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">gatherv2</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">SparseGatherV2</span><span class="p">()</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">gatherv2</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Gather</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">embeddinglookup</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">EmbeddingLookup</span><span class="p">()</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;primitive_target&#39;</span><span class="p">,</span> <span class="s1">&#39;CPU&#39;</span><span class="p">)</span>
        <span class="n">enable_ps</span> <span class="o">=</span> <span class="n">_get_ps_context</span><span class="p">(</span><span class="s2">&quot;enable_ps&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">enable_ps</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_process_vocab_cache</span><span class="p">(</span><span class="n">slice_mode</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">embedding_size</span> <span class="o">=</span> <span class="n">validator</span><span class="o">.</span><span class="n">check_positive_int</span><span class="p">(</span><span class="n">embedding_size</span><span class="p">,</span> <span class="s1">&#39;embedding_size&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">embedding_table</span> <span class="o">=</span> <span class="n">Parameter</span><span class="p">(</span><span class="n">initializer</span><span class="p">(</span><span class="n">param_init</span><span class="p">,</span> <span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">vocab_size</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">embedding_size</span><span class="p">]),</span>
                                         <span class="n">name</span><span class="o">=</span><span class="s1">&#39;embedding_table&#39;</span><span class="p">)</span>
        <span class="n">parallel_mode</span> <span class="o">=</span> <span class="n">_get_parallel_mode</span><span class="p">()</span>
        <span class="n">is_auto_parallel</span> <span class="o">=</span> <span class="n">parallel_mode</span> <span class="ow">in</span> <span class="p">(</span><span class="n">ParallelMode</span><span class="o">.</span><span class="n">SEMI_AUTO_PARALLEL</span><span class="p">,</span> <span class="n">ParallelMode</span><span class="o">.</span><span class="n">AUTO_PARALLEL</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">gather_revert</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Gather</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">reshape_first</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Reshape</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">reshape</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Reshape</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">unique</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Unique</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">shape</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Shape</span><span class="p">()</span>
        <span class="k">if</span> <span class="n">is_auto_parallel</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">unique</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Unique</span><span class="p">()</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,),))</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">cache_enable</span> <span class="ow">and</span> <span class="n">enable_ps</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_set_voacb_cache_enable_for_ps</span><span class="p">(</span><span class="n">vocab_cache_size</span><span class="p">,</span> <span class="n">embedding_size</span><span class="p">,</span> <span class="n">vocab_size</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">is_auto_parallel</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">unique</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;cache_enable&#39;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>
        <span class="n">indices_shape_size</span> <span class="o">=</span> <span class="mi">2</span>
        <span class="k">if</span> <span class="n">slice_mode</span> <span class="o">==</span> <span class="s2">&quot;field_slice&quot;</span> <span class="ow">and</span> <span class="n">is_auto_parallel</span><span class="p">:</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">manual_shapes</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="si">}</span><span class="s2">&#39;, the &#39;manual_shapes&#39; should not be none &quot;</span>
                                 <span class="sa">f</span><span class="s2">&quot;when the &#39;slice_mode&#39; is </span><span class="se">\&quot;</span><span class="s2">filed_slice</span><span class="se">\&quot;</span><span class="s2">, but got </span><span class="si">{</span><span class="n">manual_shapes</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">manual_shapes</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">):</span>
                <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="si">}</span><span class="s2">&#39;, the type of &#39;manual_shapes&#39; must be tuple(int), &quot;</span>
                                <span class="sa">f</span><span class="s2">&quot;but got </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">manual_shapes</span><span class="p">)</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">!&quot;</span><span class="p">)</span>
            <span class="k">for</span> <span class="n">dim</span> <span class="ow">in</span> <span class="n">manual_shapes</span><span class="p">:</span>
                <span class="n">validator</span><span class="o">.</span><span class="n">check_positive_int</span><span class="p">(</span><span class="n">dim</span><span class="p">,</span> <span class="s1">&#39;manual shape dim&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">gatherv2</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s2">&quot;manual_split&quot;</span><span class="p">,</span> <span class="n">manual_shapes</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">embeddinglookup</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s2">&quot;manual_split&quot;</span><span class="p">,</span> <span class="n">manual_shapes</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">gatherv2</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">),</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">get_group_size</span><span class="p">())))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">embeddinglookup</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">),</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">get_group_size</span><span class="p">())))</span>
        <span class="k">elif</span> <span class="n">slice_mode</span> <span class="o">==</span> <span class="s2">&quot;table_row_slice&quot;</span> <span class="ow">and</span> <span class="n">is_auto_parallel</span><span class="p">:</span>
            <span class="n">full_batch</span> <span class="o">=</span> <span class="n">_get_full_batch</span><span class="p">()</span>
            <span class="k">if</span> <span class="p">(</span><span class="n">target</span> <span class="o">==</span> <span class="s1">&#39;DEVICE&#39;</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">full_batch</span><span class="p">)</span> <span class="ow">or</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cache_enable</span> <span class="ow">and</span> <span class="n">enable_ps</span> <span class="ow">and</span> <span class="n">sparse</span><span class="p">):</span>
                <span class="n">indices_shape_size</span> <span class="o">=</span> <span class="mi">1</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">gather_revert</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="p">(</span><span class="n">get_group_size</span><span class="p">(),)))</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">forward_unique</span> <span class="o">=</span> <span class="kc">True</span>
            <span class="n">indices_strategy</span> <span class="o">=</span> <span class="p">(</span><span class="mi">1</span><span class="p">,)</span><span class="o">*</span><span class="n">indices_shape_size</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">gatherv2</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">),</span> <span class="n">indices_strategy</span><span class="p">))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">embeddinglookup</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">),</span> <span class="n">indices_strategy</span><span class="p">))</span>
        <span class="k">elif</span> <span class="n">slice_mode</span> <span class="o">==</span> <span class="s2">&quot;table_column_slice&quot;</span> <span class="ow">and</span> <span class="n">is_auto_parallel</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">target</span> <span class="o">==</span> <span class="s1">&#39;DEVICE&#39;</span><span class="p">:</span>
                <span class="n">indices_shape_size</span> <span class="o">=</span> <span class="mi">1</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">gather_revert</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,</span> <span class="n">get_group_size</span><span class="p">()),</span> <span class="p">(</span><span class="mi">1</span><span class="p">,)))</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">forward_unique</span> <span class="o">=</span> <span class="kc">True</span>
            <span class="n">indices_strategy</span> <span class="o">=</span> <span class="p">(</span><span class="mi">1</span><span class="p">,)</span><span class="o">*</span><span class="n">indices_shape_size</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">gatherv2</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,</span> <span class="n">get_group_size</span><span class="p">()),</span> <span class="n">indices_strategy</span><span class="p">))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">embeddinglookup</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,</span> <span class="n">get_group_size</span><span class="p">()),</span> <span class="n">indices_strategy</span><span class="p">))</span>
        <span class="k">elif</span> <span class="n">slice_mode</span> <span class="o">==</span> <span class="s2">&quot;batch_slice&quot;</span> <span class="ow">and</span> <span class="n">is_auto_parallel</span><span class="p">:</span>
            <span class="n">indices_strategy</span> <span class="o">=</span> <span class="p">[</span><span class="n">get_group_size</span><span class="p">()]</span>
            <span class="n">indices_strategy</span><span class="o">.</span><span class="n">extend</span><span class="p">([</span><span class="mi">1</span><span class="p">]</span><span class="o">*</span><span class="p">(</span><span class="n">indices_shape_size</span> <span class="o">-</span> <span class="mi">1</span><span class="p">))</span>
            <span class="n">indices_strategy</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">indices_strategy</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">gatherv2</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="n">indices_strategy</span><span class="p">))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">embeddinglookup</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="n">indices_strategy</span><span class="p">))</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">is_auto_parallel</span><span class="p">:</span>
                <span class="n">support_mode</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;field_slice&quot;</span><span class="p">,</span> <span class="s2">&quot;table_row_slice&quot;</span><span class="p">,</span> <span class="s2">&quot;table_column_slice&quot;</span><span class="p">,</span> <span class="s2">&quot;batch_slice&quot;</span><span class="p">]</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;For &#39;</span><span class="si">{}</span><span class="s2">&#39;, the &#39;slice_mode&#39; must be in </span><span class="si">{}</span><span class="s2">, &quot;</span>
                                 <span class="s2">&quot;but got </span><span class="se">\&quot;</span><span class="si">{}</span><span class="se">\&quot;</span><span class="s2">.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">,</span> <span class="n">support_mode</span><span class="p">,</span> <span class="n">slice_mode</span><span class="p">))</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">cache_enable</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">enable_ps</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="si">}</span><span class="s2">&#39;, haven&#39;t supported cache enable for not ps mode.&quot;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">embedding_table</span><span class="o">.</span><span class="n">unique</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">forward_unique</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">max_norm</span> <span class="o">=</span> <span class="n">max_norm</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">max_norm</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">max_norm</span> <span class="o">=</span> <span class="n">validator</span><span class="o">.</span><span class="n">check_positive_float</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">max_norm</span><span class="p">,</span> <span class="s1">&#39;max_norm&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">max_norm</span> <span class="o">=</span> <span class="n">Tensor</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">max_norm</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_process_vocab_cache</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">slice_mode</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;PS embeddingLookup cache check and process.&quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cache_enable</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">vocab_cache_size</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">target</span> <span class="o">==</span> <span class="s1">&#39;CPU&#39;</span><span class="p">:</span>
                <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;The configuration of &#39;vocab_cache_size&#39; is valid only in &#39;DEVICE&#39; target, &quot;</span>
                               <span class="s2">&quot;current target is CPU, so it will be ignored.&quot;</span><span class="p">)</span>
                <span class="k">return</span>
            <span class="n">enable_ps</span> <span class="o">=</span> <span class="n">_get_ps_context</span><span class="p">(</span><span class="s2">&quot;enable_ps&quot;</span><span class="p">)</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">enable_ps</span><span class="p">:</span>
                <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;The configuration of &#39;vocab_cache_size&#39; is valid only in parameter server training &quot;</span>
                               <span class="s2">&quot;mode, current mode is not parameter server trainning mode, so it will be ignored.&quot;</span><span class="p">)</span>
                <span class="k">return</span>
            <span class="n">parallel_mode</span> <span class="o">=</span> <span class="n">_get_parallel_mode</span><span class="p">()</span>
            <span class="n">is_auto_parallel</span> <span class="o">=</span> <span class="n">parallel_mode</span> <span class="ow">in</span> <span class="p">(</span><span class="n">ParallelMode</span><span class="o">.</span><span class="n">SEMI_AUTO_PARALLEL</span><span class="p">,</span> <span class="n">ParallelMode</span><span class="o">.</span><span class="n">AUTO_PARALLEL</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">is_auto_parallel</span><span class="p">:</span>
                <span class="n">rank_size</span> <span class="o">=</span> <span class="n">get_group_size</span><span class="p">()</span>
                <span class="n">rank_id</span> <span class="o">=</span> <span class="n">get_rank</span><span class="p">()</span>
                <span class="n">full_batch</span> <span class="o">=</span> <span class="n">_get_full_batch</span><span class="p">()</span>
                <span class="k">if</span> <span class="n">rank_size</span> <span class="o">&gt;</span> <span class="mi">1</span> <span class="ow">and</span> <span class="ow">not</span> <span class="p">(</span><span class="n">full_batch</span> <span class="ow">and</span> <span class="n">slice_mode</span> <span class="o">==</span> <span class="s2">&quot;table_row_slice&quot;</span><span class="p">):</span>
                    <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="si">}</span><span class="s2">&#39;, the cache of parameter server parallel should only be &quot;</span>
                                     <span class="sa">f</span><span class="s2">&quot;used in </span><span class="se">\&quot;</span><span class="s2">full_batch</span><span class="se">\&quot;</span><span class="s2"> and the value of </span><span class="se">\&quot;</span><span class="s2">full_batch</span><span class="se">\&quot;</span><span class="s2"> should be True. &quot;</span>
                                     <span class="sa">f</span><span class="s2">&quot;Meanwhile, the value of &#39;slice_mode&#39; should be </span><span class="se">\&quot;</span><span class="s2">table_row_slice</span><span class="se">\&quot;</span><span class="s2">.&quot;</span>
                                     <span class="sa">f</span><span class="s2">&quot;But got full_batch: </span><span class="si">{</span><span class="n">full_batch</span><span class="si">}</span><span class="s2"> and &#39;slice_mode&#39;: </span><span class="se">\&quot;</span><span class="si">{</span><span class="n">slice_mode</span><span class="si">}</span><span class="se">\&quot;</span><span class="s2">.&quot;</span><span class="p">)</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">vocab_cache_size</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">vocab_cache_size</span> <span class="o">*</span> <span class="n">rank_size</span>
                <span class="n">_set_rank_id</span><span class="p">(</span><span class="n">rank_id</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">cache_enable</span> <span class="o">=</span> <span class="kc">True</span>
            <span class="k">if</span> <span class="n">_is_role_worker</span><span class="p">():</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">vocab_size</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">vocab_cache_size</span>
                <span class="k">if</span> <span class="n">context</span><span class="o">.</span><span class="n">get_context</span><span class="p">(</span><span class="s2">&quot;enable_sparse&quot;</span><span class="p">)</span> <span class="o">!=</span> <span class="bp">self</span><span class="o">.</span><span class="n">sparse</span><span class="p">:</span>
                    <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="si">}</span><span class="s2">&#39;, the value of parameter &#39;sparse&#39; must be same for all &quot;</span>
                                     <span class="sa">f</span><span class="s2">&quot;kernels and equal the value of &#39;enable_sparse&#39; in context setting in &quot;</span>
                                     <span class="sa">f</span><span class="s2">&quot;parameter server cache mode, but got value of parameter &#39;sparse&#39;: </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">sparse</span><span class="si">}</span><span class="s2">&quot;</span>
                                     <span class="sa">f</span><span class="s2">&quot; and the &#39;enable_sparse&#39; in context setting: &quot;</span>
                                     <span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">context</span><span class="o">.</span><span class="n">get_context</span><span class="p">(</span><span class="s1">&#39;enable_sparse&#39;</span><span class="p">)</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_set_voacb_cache_enable_for_ps</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">vocab_cache_size</span><span class="p">,</span> <span class="n">embedding_size</span><span class="p">,</span> <span class="n">vocab_size</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;PS embeddingLookup cache enable set.&quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">embedding_table</span><span class="o">.</span><span class="n">cache_enable</span> <span class="o">=</span> <span class="kc">True</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">embedding_table</span><span class="o">.</span><span class="n">is_param_ps</span> <span class="o">=</span> <span class="kc">True</span>
        <span class="n">_set_cache_enable</span><span class="p">(</span><span class="kc">True</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">sparse</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">forward_unique</span> <span class="o">=</span> <span class="kc">True</span>
        <span class="k">if</span> <span class="n">_is_role_worker</span><span class="p">():</span>
            <span class="n">_insert_hash_table_size</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">embedding_table</span><span class="o">.</span><span class="n">name</span><span class="p">,</span> <span class="n">vocab_cache_size</span><span class="p">,</span> <span class="n">embedding_size</span><span class="p">,</span> <span class="n">vocab_size</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">construct</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">indices</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">target</span> <span class="o">==</span> <span class="s2">&quot;CPU&quot;</span><span class="p">:</span>
            <span class="n">out</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">embeddinglookup</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">embedding_table</span><span class="p">,</span> <span class="n">indices</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">forward_unique</span><span class="p">:</span>
                <span class="n">shp</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">indices</span><span class="p">)</span> <span class="o">+</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">embedding_size</span><span class="p">,)</span>
                <span class="n">indices_flatten</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshape_first</span><span class="p">(</span><span class="n">indices</span><span class="p">,</span> <span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,))</span>
                <span class="n">unique_id</span><span class="p">,</span> <span class="n">unique_idx</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">unique</span><span class="p">(</span><span class="n">indices_flatten</span><span class="p">)</span>
                <span class="n">weight_unique</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">gatherv2</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">embedding_table</span><span class="p">,</span> <span class="n">unique_id</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
                <span class="n">weight_flatten</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">gather_revert</span><span class="p">(</span><span class="n">weight_unique</span><span class="p">,</span> <span class="n">unique_idx</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
                <span class="n">out</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">weight_flatten</span><span class="p">,</span> <span class="n">shp</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">out</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">gatherv2</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">embedding_table</span><span class="p">,</span> <span class="n">indices</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">max_norm</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">axis</span> <span class="o">=</span> <span class="n">_make_axis_range</span><span class="p">(</span><span class="n">F</span><span class="o">.</span><span class="n">rank</span><span class="p">(</span><span class="n">indices</span><span class="p">),</span> <span class="n">F</span><span class="o">.</span><span class="n">rank</span><span class="p">(</span><span class="n">out</span><span class="p">))</span>
            <span class="n">clip_by_norm</span> <span class="o">=</span> <span class="n">ClipByNorm</span><span class="p">(</span><span class="n">axis</span><span class="p">)</span>
            <span class="n">out</span> <span class="o">=</span> <span class="n">clip_by_norm</span><span class="p">(</span><span class="n">out</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">max_norm</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">out</span>


<div class="viewcode-block" id="MultiFieldEmbeddingLookup"><a class="viewcode-back" href="../../../../api_python/nn/mindspore.nn.MultiFieldEmbeddingLookup.html#mindspore.nn.MultiFieldEmbeddingLookup">[docs]</a><span class="k">class</span> <span class="nc">MultiFieldEmbeddingLookup</span><span class="p">(</span><span class="n">EmbeddingLookup</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Returns a slice of input tensor based on the specified indices and the field ids. This operation</span>
<span class="sd">    supports looking up embeddings using multi hot and one hot fields simultaneously.</span>

<span class="sd">    Note:</span>
<span class="sd">        When &#39;target&#39; is set to &#39;CPU&#39;, this module will use</span>
<span class="sd">        P.EmbeddingLookup().add_prim_attr(&#39;primitive_target&#39;, &#39;CPU&#39;) which</span>
<span class="sd">        specified &#39;offset = 0&#39; to lookup table.</span>
<span class="sd">        When &#39;target&#39; is set to &#39;DEVICE&#39;, this module will use P.Gather() which</span>
<span class="sd">        specified &#39;axis = 0&#39; to lookup table.</span>
<span class="sd">        The vectors with the same field_ids  will be combined by the `operator`, such as &#39;SUM&#39;, &#39;MAX&#39; and</span>
<span class="sd">        &#39;MEAN&#39;. Ensure the input_values of the padded id is zero, so that they can be ignored. The final</span>
<span class="sd">        output will be zeros if the sum of absolute weight of the field is zero. This class only</span>
<span class="sd">        supports [&#39;table_row_slice&#39;, &#39;batch_slice&#39; and &#39;table_column_slice&#39;]. For the operation &#39;MAX&#39; on</span>
<span class="sd">        device Ascend, there is a constraint where batch_size * (seq_length + field_size) &lt; 3500.</span>

<span class="sd">    Args:</span>
<span class="sd">        vocab_size (int): The size of the dictionary of embeddings.</span>
<span class="sd">        embedding_size (int): The size of each embedding vector.</span>
<span class="sd">        field_size (int): The field size of the final outputs.</span>
<span class="sd">        param_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the embedding_table.</span>
<span class="sd">            Refer to class `initializer` for the values of string when a string</span>
<span class="sd">            is specified. Default: &#39;normal&#39;.</span>
<span class="sd">        target (str): Specifies the target where the op is executed. The value must in</span>
<span class="sd">            [&#39;DEVICE&#39;, &#39;CPU&#39;]. Default: &#39;CPU&#39;.</span>
<span class="sd">        slice_mode (str): The slicing way in semi_auto_parallel/auto_parallel. The value must get through</span>
<span class="sd">            :class:`mindspore.nn.EmbeddingLookup`. Default: :class:`mindspore.nn.EmbeddingLookup.BATCH_SLICE`.</span>
<span class="sd">        feature_num_list (tuple): The accompaniment array in field slice mode. This is unused currently. Default: None.</span>
<span class="sd">        max_norm (Union[float, None]): A maximum clipping value. The data type must be float16, float32</span>
<span class="sd">                                       or None. Default: None</span>
<span class="sd">        sparse (bool): Using sparse mode. When &#39;target&#39; is set to &#39;CPU&#39;, &#39;sparse&#39; has to be true. Default: True.</span>
<span class="sd">        operator (str): The pooling method for the features in one field. Support &#39;SUM&#39;, &#39;MEAN&#39; and &#39;MAX&#39;.</span>
<span class="sd">            Default: &#39;SUM&#39;.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **input_indices** (Tensor) - The shape of tensor is :math:`(batch\_size, seq\_length)`.</span>
<span class="sd">          Specifies the indices of elements of the original Tensor. Input_indices must be a 2d tensor in</span>
<span class="sd">          this interface. Type is Int32, Int64.</span>
<span class="sd">        - **input_values** (Tensor) - The shape of tensor is :math:`(batch\_size, seq\_length)`.</span>
<span class="sd">          Specifies the weights of elements of the input_indices. The lookout vector will multiply with</span>
<span class="sd">          the input_values. Type is Float32.</span>
<span class="sd">        - **field_ids** (Tensor)  - The shape of tensor is :math:`(batch\_size, seq\_length)`.</span>
<span class="sd">          Specifies the field id of elements of the input_indices. Type is Int32.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor, the shape of tensor is :math:`(batch\_size, field\_size, embedding\_size)`. Type is Float32.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If `vocab_size` or `embedding_size` or `field_size` is not an int.</span>
<span class="sd">        TypeError: If `sparse` is not a bool or `feature_num_list` is not a tuple.</span>
<span class="sd">        ValueError: If `vocab_size` or `embedding_size` or `field_size` is less than 1.</span>
<span class="sd">        ValueError: If `target` is neither &#39;CPU&#39; nor &#39;DEVICE&#39;.</span>
<span class="sd">        ValueError: If `slice_mode` is not one of &#39;batch_slice&#39;, &#39;field_slice&#39;, &#39;table_row_slice&#39;,</span>
<span class="sd">                    &#39;table_column_slice&#39;.</span>
<span class="sd">        ValueError: If `sparse` is False and `target` is &#39;CPU&#39;.</span>
<span class="sd">        ValueError: If `slice_mode` is &#39;field_slice&#39; and `feature_num_list` is None.</span>
<span class="sd">        ValueError: If `operator` is not one of &#39;SUM&#39;, &#39;MAX&#39;, &#39;MEAN&#39;.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; input_indices = Tensor([[2, 4, 6, 0, 0], [1, 3, 5, 0, 0]], mindspore.int32)</span>
<span class="sd">        &gt;&gt;&gt; input_values = Tensor([[1, 1, 1, 0, 0], [1, 1, 1, 0, 0]], mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; field_ids = Tensor([[0, 1, 1, 0, 0], [0, 0, 1, 0, 0]], mindspore.int32)</span>
<span class="sd">        &gt;&gt;&gt; net = nn.MultiFieldEmbeddingLookup(10, 2, field_size=2, operator=&#39;SUM&#39;, target=&#39;DEVICE&#39;)</span>
<span class="sd">        &gt;&gt;&gt; out = net(input_indices, input_values, field_ids)</span>
<span class="sd">        &gt;&gt;&gt; print(out.shape)</span>
<span class="sd">        (2, 2, 2)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">OPERATOR_SUM</span> <span class="o">=</span> <span class="s1">&#39;SUM&#39;</span>
    <span class="n">OPERATOR_MEAN</span> <span class="o">=</span> <span class="s1">&#39;MEAN&#39;</span>
    <span class="n">OPERATOR_MAX</span> <span class="o">=</span> <span class="s1">&#39;MAX&#39;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">vocab_size</span><span class="p">,</span> <span class="n">embedding_size</span><span class="p">,</span> <span class="n">field_size</span><span class="p">,</span> <span class="n">param_init</span><span class="o">=</span><span class="s1">&#39;normal&#39;</span><span class="p">,</span> <span class="n">target</span><span class="o">=</span><span class="s1">&#39;CPU&#39;</span><span class="p">,</span>
                 <span class="n">slice_mode</span><span class="o">=</span><span class="s1">&#39;batch_slice&#39;</span><span class="p">,</span> <span class="n">feature_num_list</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">max_norm</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">sparse</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">operator</span><span class="o">=</span><span class="s1">&#39;SUM&#39;</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize MultiFieldEmbeddingLookup.&quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">MultiFieldEmbeddingLookup</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">vocab_size</span><span class="p">,</span> <span class="n">embedding_size</span><span class="p">,</span> <span class="n">param_init</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span>
                                                        <span class="n">slice_mode</span><span class="p">,</span> <span class="n">feature_num_list</span><span class="p">,</span> <span class="n">max_norm</span><span class="p">,</span> <span class="n">sparse</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">field_size</span> <span class="o">=</span> <span class="n">validator</span><span class="o">.</span><span class="n">check_positive_int</span><span class="p">(</span><span class="n">field_size</span><span class="p">,</span> <span class="s1">&#39;field_size&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">operator</span> <span class="o">=</span> <span class="n">operator</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">mul</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Mul</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">inf_mask_mul</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Mul</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">bias_add</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Add</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">inf_add</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Add</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">merge_op</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">count_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">UnsortedSegmentSum</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">abs</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Abs</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">equal</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Equal</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Add</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cast</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Cast</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">div_no_nan</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">DivNoNan</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">expand</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">ExpandDims</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">max_mask_mul</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Mul</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">max_no_equal</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">NotEqual</span><span class="p">()</span>

        <span class="n">validator</span><span class="o">.</span><span class="n">check_string</span><span class="p">(</span><span class="n">operator</span><span class="p">,</span> <span class="p">[</span><span class="s1">&#39;SUM&#39;</span><span class="p">,</span> <span class="s1">&#39;MAX&#39;</span><span class="p">,</span> <span class="s1">&#39;MEAN&#39;</span><span class="p">],</span> <span class="s1">&#39;operator&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">operator</span> <span class="o">==</span> <span class="n">MultiFieldEmbeddingLookup</span><span class="o">.</span><span class="n">OPERATOR_SUM</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">merge_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">UnsortedSegmentSum</span><span class="p">()</span>
        <span class="k">elif</span> <span class="n">operator</span> <span class="o">==</span> <span class="n">MultiFieldEmbeddingLookup</span><span class="o">.</span><span class="n">OPERATOR_MAX</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">merge_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">UnsortedSegmentMax</span><span class="p">()</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">merge_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">UnsortedSegmentSum</span><span class="p">()</span>


        <span class="n">parallel_mode</span> <span class="o">=</span> <span class="n">_get_parallel_mode</span><span class="p">()</span>
        <span class="n">is_auto_parallel</span> <span class="o">=</span> <span class="n">parallel_mode</span> <span class="ow">in</span> <span class="p">(</span><span class="n">ParallelMode</span><span class="o">.</span><span class="n">SEMI_AUTO_PARALLEL</span><span class="p">,</span> <span class="n">ParallelMode</span><span class="o">.</span><span class="n">AUTO_PARALLEL</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">slice_mode</span> <span class="ow">in</span> <span class="p">[</span><span class="s2">&quot;table_row_slice&quot;</span><span class="p">,</span> <span class="s2">&quot;batch_slice&quot;</span><span class="p">]</span> <span class="ow">and</span> <span class="n">is_auto_parallel</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">merge_op</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="p">(</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">)))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">expand</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="n">get_group_size</span><span class="p">(),),))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">bias_add</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">mul</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="p">(</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">count_op</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">),</span> <span class="p">(</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">)))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">add</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="n">get_group_size</span><span class="p">(),),</span> <span class="p">(</span><span class="n">get_group_size</span><span class="p">(),)))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">div_no_nan</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">),</span> <span class="p">(</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">)))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">max_mask_mul</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">),</span> <span class="p">(</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">)))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">max_no_equal</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,),</span> <span class="p">()))</span>
            <span class="k">if</span> <span class="n">operator</span> <span class="o">==</span> <span class="n">MultiFieldEmbeddingLookup</span><span class="o">.</span><span class="n">OPERATOR_MAX</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">equal</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="p">()))</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">inf_mask_mul</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="p">()))</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">merge_op</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">),</span> <span class="p">(</span><span class="n">get_group_size</span><span class="p">(),)))</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">count_op</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="n">get_group_size</span><span class="p">(),),</span> <span class="p">(</span><span class="n">get_group_size</span><span class="p">(),)))</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">inf_add</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="p">(</span><span class="n">get_group_size</span><span class="p">(),</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)))</span>
        <span class="k">elif</span> <span class="n">slice_mode</span> <span class="o">==</span> <span class="s2">&quot;table_column_slice&quot;</span> <span class="ow">and</span> <span class="n">is_auto_parallel</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">merge_op</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">get_group_size</span><span class="p">()),</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">div_no_nan</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,</span> <span class="n">get_group_size</span><span class="p">()),</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">bias_add</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">mul</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">get_group_size</span><span class="p">())))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">count_op</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">add</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,),</span> <span class="p">(</span><span class="mi">1</span><span class="p">,)))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">max_mask_mul</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,</span> <span class="n">get_group_size</span><span class="p">()),</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">expand</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,),))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">max_no_equal</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,),</span> <span class="p">()))</span>
            <span class="k">if</span> <span class="n">operator</span> <span class="o">==</span> <span class="n">MultiFieldEmbeddingLookup</span><span class="o">.</span><span class="n">OPERATOR_MAX</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">equal</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="p">()))</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">inf_mask_mul</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="p">()))</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">merge_op</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,</span> <span class="n">get_group_size</span><span class="p">()),</span> <span class="p">(</span><span class="mi">1</span><span class="p">,)))</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">count_op</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,),</span> <span class="p">(</span><span class="mi">1</span><span class="p">,)))</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">inf_add</span><span class="o">.</span><span class="n">shard</span><span class="p">(((</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">get_group_size</span><span class="p">()),</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)))</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">is_auto_parallel</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;For &#39;</span><span class="si">{}</span><span class="s2">&#39;, the &#39;slice_mode&#39; should be in [&#39;table_row_slice&#39;, &#39;batch_slice&#39; and </span><span class="se">\</span>
<span class="s2">                                       &#39;table_column_slice&#39;], but got </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">,</span> <span class="nb">str</span><span class="p">(</span><span class="n">slice_mode</span><span class="p">)))</span>

        <span class="c1"># Min value for fp32</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">negative_inf_value</span> <span class="o">=</span> <span class="o">-</span><span class="mf">3.402823466E+38</span>

    <span class="k">def</span> <span class="nf">construct</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_indices</span><span class="p">,</span> <span class="n">input_values</span><span class="p">,</span> <span class="n">field_ids</span><span class="p">):</span>

        <span class="n">_check_input_2d</span><span class="p">(</span><span class="n">F</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">input_indices</span><span class="p">),</span> <span class="s2">&quot;input_indices&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="n">_check_input_2d</span><span class="p">(</span><span class="n">F</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">input_values</span><span class="p">),</span> <span class="s2">&quot;input_values&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="n">_check_input_2d</span><span class="p">(</span><span class="n">F</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">field_ids</span><span class="p">),</span> <span class="s2">&quot;field_ids&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="n">_check_input_dtype</span><span class="p">(</span><span class="n">F</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">input_indices</span><span class="p">),</span> <span class="s2">&quot;input_indices&quot;</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">int32</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">int64</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="n">_check_input_dtype</span><span class="p">(</span><span class="n">F</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">input_values</span><span class="p">),</span> <span class="s2">&quot;input_values&quot;</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="n">_check_input_dtype</span><span class="p">(</span><span class="n">F</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">field_ids</span><span class="p">),</span> <span class="s2">&quot;field_ids&quot;</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">int32</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>

        <span class="n">batch_size</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">input_indices</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span>
        <span class="n">num_segments</span> <span class="o">=</span> <span class="n">batch_size</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">field_size</span>
        <span class="n">bias</span> <span class="o">=</span> <span class="n">Range</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">num_segments</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">field_size</span><span class="p">)()</span>
        <span class="n">bias</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">bias</span><span class="p">,</span> <span class="p">(</span><span class="n">batch_size</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">))</span>
        <span class="n">field_ids</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">bias_add</span><span class="p">(</span><span class="n">field_ids</span><span class="p">,</span> <span class="n">bias</span><span class="p">)</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">target</span> <span class="o">==</span> <span class="s2">&quot;CPU&quot;</span><span class="p">:</span>
            <span class="n">out</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">embeddinglookup</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">embedding_table</span><span class="p">,</span> <span class="n">input_indices</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">forward_unique</span><span class="p">:</span>
                <span class="n">shp</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">input_indices</span><span class="p">)</span> <span class="o">+</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">embedding_size</span><span class="p">,)</span>
                <span class="n">indices_flatten</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">input_indices</span><span class="p">,</span> <span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,))</span>
                <span class="n">unique_id</span><span class="p">,</span> <span class="n">unique_idx</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">unique</span><span class="p">(</span><span class="n">indices_flatten</span><span class="p">)</span>
                <span class="n">weight_unique</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">gatherv2</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">embedding_table</span><span class="p">,</span> <span class="n">unique_id</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
                <span class="n">weight_flatten</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">gather_revert</span><span class="p">(</span><span class="n">weight_unique</span><span class="p">,</span> <span class="n">unique_idx</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
                <span class="n">out</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">weight_flatten</span><span class="p">,</span> <span class="n">shp</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">out</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">gatherv2</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">embedding_table</span><span class="p">,</span> <span class="n">input_indices</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">max_norm</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">axis</span> <span class="o">=</span> <span class="n">_make_axis_range</span><span class="p">(</span><span class="n">F</span><span class="o">.</span><span class="n">rank</span><span class="p">(</span><span class="n">input_indices</span><span class="p">),</span> <span class="n">F</span><span class="o">.</span><span class="n">rank</span><span class="p">(</span><span class="n">out</span><span class="p">))</span>
            <span class="n">clip_by_norm</span> <span class="o">=</span> <span class="n">ClipByNorm</span><span class="p">(</span><span class="n">axis</span><span class="p">)</span>
            <span class="n">out</span> <span class="o">=</span> <span class="n">clip_by_norm</span><span class="p">(</span><span class="n">out</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">max_norm</span><span class="p">)</span>

        <span class="n">weights</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">input_values</span><span class="p">,</span> <span class="p">(</span><span class="n">batch_size</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">input_indices</span><span class="p">)[</span><span class="mi">1</span><span class="p">],</span> <span class="mi">1</span><span class="p">))</span>
        <span class="n">embedding</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">mul</span><span class="p">(</span><span class="n">weights</span><span class="p">,</span> <span class="n">out</span><span class="p">)</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">operator</span> <span class="o">==</span> <span class="s1">&#39;MAX&#39;</span><span class="p">:</span>
            <span class="c1"># Fill the padding value to -inf, so the padded value will not influence the results</span>
            <span class="n">negative_inf_mask</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">cast</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">equal</span><span class="p">(</span><span class="n">weights</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>
            <span class="n">inf_mask</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">inf_mask_mul</span><span class="p">(</span><span class="n">negative_inf_mask</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">negative_inf_value</span><span class="p">)</span>
            <span class="n">embedding</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">inf_add</span><span class="p">(</span><span class="n">embedding</span><span class="p">,</span> <span class="n">inf_mask</span><span class="p">)</span>
            <span class="n">embedding</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">embedding</span><span class="p">,</span> <span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">embedding_size</span><span class="p">))</span>
            <span class="n">field_ids</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">field_ids</span><span class="p">,</span> <span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,))</span>

        <span class="n">merged_vectors</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">merge_op</span><span class="p">(</span><span class="n">embedding</span><span class="p">,</span> <span class="n">field_ids</span><span class="p">,</span> <span class="n">num_segments</span><span class="p">)</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">operator</span> <span class="o">==</span> <span class="s1">&#39;MAX&#39;</span><span class="p">:</span>
            <span class="n">value_count</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">count_op</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">input_values</span><span class="p">,</span> <span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,))),</span> <span class="n">field_ids</span><span class="p">,</span> <span class="n">num_segments</span><span class="p">)</span>
            <span class="n">value_zeros</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">cast</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">max_no_equal</span><span class="p">(</span><span class="n">value_count</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">),</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>
            <span class="n">count</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">expand</span><span class="p">(</span><span class="n">value_zeros</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>
            <span class="n">merged_vectors</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">max_mask_mul</span><span class="p">(</span><span class="n">merged_vectors</span><span class="p">,</span> <span class="n">count</span><span class="p">)</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">operator</span> <span class="o">==</span> <span class="s1">&#39;MEAN&#39;</span><span class="p">:</span>
            <span class="n">value_count</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">count_op</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="n">input_values</span><span class="p">),</span> <span class="n">field_ids</span><span class="p">,</span> <span class="n">num_segments</span><span class="p">)</span>
            <span class="n">value_count</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">expand</span><span class="p">(</span><span class="n">value_count</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>
            <span class="n">merged_vectors</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">div_no_nan</span><span class="p">(</span><span class="n">merged_vectors</span><span class="p">,</span> <span class="n">value_count</span><span class="p">)</span>

        <span class="n">merged_vectors</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">merged_vectors</span><span class="p">,</span> <span class="p">(</span><span class="n">batch_size</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">field_size</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">))</span>
        <span class="k">return</span> <span class="n">merged_vectors</span></div>
</pre></div>

           </div>
           
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>
        &#169; Copyright 2021, MindSpore.

    </p>
  </div>
    
    
    
    Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    
    provided by <a href="https://readthedocs.org">Read the Docs</a>. 

</footer>
        </div>
      </div>

    </section>

  </div>
  

  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>

  
  
    
   

</body>
</html>